%clear; %clc; close all;
rng('default');

%%% Set Parameters
periods_within_year = 1;


sigma_l = 0.2/sqrt(periods_within_year); % std dev of idio iid endowment shocks
lbar    =  - 0.5*sigma_l^2;
abar    = 0;
p_unemp   = 0;
w_unemp = 0.05;

% households
gam = 1; % CRRA
b = 0; %-0.75*1.2; % borrowing constraint
deltax = 0.02/periods_within_year; % death probability
delta_eta = 0;
delta_inh = 1;
ex_signals = 0*1.3065;

bta = 0.96^(1/periods_within_year); % discount factor
btat = bta*(1-deltax*(1-delta_inh)); % subjective discount factor (includes death)

W_cc = 1;
kappa =  0.9656; 1.075;     %/3;
sigma_c =  sqrt(0.7658);  sqrt(0.9689);  sqrt(0.9); sqrt(0.8049);
psi =    0.05;  
cov_fun = 1;
factor_mpc1 = 1;    %% in terms of standard deviations
factor_mpc2 = 10;   %%%% in terms of standard deviations
sigma_tremble = 0.183;
 
[kappa, sigma_c^2, psi, b]

% firms
alpha = 0.36; % capital share
delta = 0.08/periods_within_year; % depreciation rate
Z     = 1;

util_type = 'crra';
run_fsolve   = 0;
log_c = 0;


%%% Finds equilibrium interest rate following algorithm similar to Violante notes

% Stationary distribution parameters and preallocation
N             = 3;               % set number of agents in cross section to simulate for stat. distr.
T             = 250;
maxiter_distr = T;   % max number of iterations in search for stationary distr.

%%%%% Big simulation
sim_periods = maxiter_distr;
nsim = 3; N;



yi = exp(normrnd( lbar,sigma_l,sim_periods,2)); % preallocate all y-shocks
unemp_shocks = rand(sim_periods, 2) < p_unemp;
yi(unemp_shocks == 1) = w_unemp*mean(yi(:));

eps_etai = normrnd(0,1, sim_periods, 2); % preallocate eta shocks
reset_shocks = ones(sim_periods, 3); rand(sim_periods, 3); % reset-shocks
eps_etai(:,3) = normrnd(0,1,sim_periods,1); 
yi(:,3)       = exp(normrnd( lbar,sigma_l,sim_periods,1));

%%%% Make agents 1 and 2 the same
%eps_etai(:,2) = eps_etai(:,1);
eps_etai(1,1) = 0.5; % agent 1 repeatedly low C
eps_etai(1,2) = -0.5; % agent 2 repeatedly high C
eps_etai(1,3) = -0.4; 
eps_etai(9:10,3) = 0.1*eps_etai(9:10,3); 
%eps_etai(6:7,1) = 0; 
%eps_etai([35,45],1) = 0.31;
eps_etai(10,1) = 0; 
eps_etai(9,2) = -1; 
%eps_etai(10,2) = 0; 
%eps_etai(9,3) = -1; 
% 
yi(:,2) = yi(:,1);
yi(:,3) = yi(:,1); 
yi(1,2) = yi(1,2) + 2;
yi(1,3) = yi(1,3) + 5;

% yi(1,2) =  0.5;
% yi(2,2) =  2;
% yi(:,2) = yi(:,1);
% yi(1,2) = yi(1,2) + 6;



% Parameters for interest rate search (bisection method)
% bounds for interst rate (be careful with these)
rLower =  0.035; -delta;
rUpper =  btat.^-1 - 1; % - 0.9999;

r0 = mean([rLower rUpper]); % initial guess

conv_r = 1;
counter_r = 0;
maxiter_r = 100;             % max number of iterations (search for r)
crit_r = 1e-5;              % stopping criteria when abs(r-rnew)<crit_r  (search for r)
crit_distr = 1e-1;      % convergence criteria for stationary distr.

amin = b;
amax = 20;
na = 101;

stateGrid_size = 1001;

%%% keep track of interest rate and market clearing in each iteration

param_fp = [bta,r0,abar,W_cc,kappa,deltax,sigma_c,psi,cov_fun,gam,sigma_l, lbar, delta_eta, delta_inh, delta, alpha,b, btat, amin, amax, na, crit_distr];

[int_nodes, int_weights] = lgwt(101,-10,10);
normpdf_nodes            = normpdf(int_nodes,0,1);
%%%
%%%%% ENDOGENOUS SIGNALS MODEL

%%%  Find the endogenous signals r

param_fp(4) = gam; %%% make the Wcc entry equal to risk-aversion, so that we do not have to redefine inputs to the update file

rt_br = 0.0389;
r_br  = rt_br;
w_br = 1.1936;


[xs,~,cstar,~,ctilde] = bc_opt_pol_iid_egm(btat, rt_br, b, log(w_br) + lbar, sigma_l, gam, amin, amax, na, util_type, p_unemp, w_unemp);
[xs,xi] = unique(xs);


% get policy function over an equally spaced grid
stateGrid_size = 1001;
stateGrid = linspace( min(xs(:)), max(xs(:)),stateGrid_size);
stateGrid_step = stateGrid(2) - stateGrid(1);

% ctilde is the optimal choice for consumption if the borrowing
% constraint were relaxed only in the current period.
if log_c == 0
    temp = griddedInterpolant(xs,ctilde(xi),'linear','linear');
else
    temp = griddedInterpolant(xs,log(ctilde(xi)),'linear','linear');
end

mu0_fun = temp;
ctilde = temp(stateGrid);

temp = griddedInterpolant(xs,cstar(xi),'linear','linear');

cstar_fixedpoints = temp(stateGrid);
temp = griddedInterpolant(stateGrid,cstar_fixedpoints,'linear','linear');
%cstar = @(x) temp(x);
cstar = @(x) fast_interp( x, cstar_fixedpoints, stateGrid_step, stateGrid,stateGrid_size);
cstar_brsim = mu0_fun;

rw_pol = (rt_br/(1+rt_br))*stateGrid' + 1/(1+rt_br)*w_br;
rw_pol_interp = griddedInterpolant(stateGrid, rw_pol,'linear','linear');


%{

plot( stateGrid, -gam*(cstar(stateGrid)').^(-gam-1))
plot( stateGrid, [ctilde', rw_pol])
xlim( [0,5])


%}


param_br_sim = [bta,rt_br,abar,W_cc,kappa,deltax,sigma_c,psi,cov_fun,gam,sigma_l, lbar, delta_eta, delta_inh, factor_mpc1, factor_mpc2,0];
param_br_sim(4) = gam;

%tic
%[a_sims_br,chat_sims_br,~,aphat_sims_br,~,eta_sims,sigma_etasq_sims,alpha_star_sims,mpc,age_sims_br,prior_sims] = bc_endsignals_sim_fast(w_br,param_br_sim,yi(1:sim_periods,1:nsim),eps_etai(1:sim_periods,1:nsim),reset_shocks(1:sim_periods,1:nsim),b,mu0_fun, ctilde, stateGrid,stateGrid_step,stateGrid_size);
%toc

tic
[a_sims_br,chat_sims_br, cstar_sims_br, aphat_sims_br,astar_sims_br,eta_sims,sigma_etasq_sims,alpha_star_sims,mpc_br,mpc_br2, mpc_br3, mpc_br4, age_sims_br,prior_sims, chat_prime_sims, alpha_mpc_sims] = bc_endsignals_sim(w_br,param_br_sim,yi(1:sim_periods,1:nsim),eps_etai(1:sim_periods,1:nsim),reset_shocks(1:sim_periods,1:nsim),b,mu0_fun, ctilde, stateGrid,stateGrid_step,stateGrid_size,cstar,log_c, abar*ones(1,nsim), ex_signals);
toc

ya = yi*w_br + a_sims_br*(1+rt_br);

%%%
%%%
tt = 1;

i=1;
ind_temp1 = 1:tt;
ind_temp1 = ind_temp1( sigma_etasq_sims(1:tt,i) < 1e10);

eta_temp = eta_sims( ind_temp1,i);

[chat_pol11,shat_pol11] =  gp_update( stateGrid' , yi(ind_temp1,i)*w_br + a_sims_br(ind_temp1,i)*(1+r_br), eta_temp, sigma_etasq_sims(ind_temp1,i), cstar_brsim,sigma_c,psi, cov_fun);

i = 2;

ind_temp2 = 1:tt;
ind_temp2 = ind_temp2( sigma_etasq_sims(1:tt,i) < 1e10);
eta_temp = eta_sims( ind_temp2,i);
[chat_pol12,shat_pol12] =  gp_update( stateGrid' , yi(ind_temp2,i)*w_br + a_sims_br(ind_temp2,i)*(1+r_br), eta_temp, sigma_etasq_sims(ind_temp2,i), cstar_brsim,sigma_c,psi, cov_fun);


tt = 100;

i=1;
ind_temp1 = 1:tt;
ind_temp1 = ind_temp1( sigma_etasq_sims(1:tt,i) < 1e10);

eta_temp = eta_sims( ind_temp1,i);

[chat_pol1,shat_pol1] =  gp_update( stateGrid' , yi(ind_temp1,i)*w_br + a_sims_br(ind_temp1,i)*(1+r_br), eta_temp, sigma_etasq_sims(ind_temp1,i), cstar_brsim,sigma_c,psi, cov_fun);

i = 2;

ind_temp2 = 1:tt;
ind_temp2 = ind_temp2( sigma_etasq_sims(1:tt,i) < 1e10);
eta_temp = eta_sims( ind_temp2,i);
[chat_pol2,shat_pol2] =  gp_update( stateGrid' , yi(ind_temp2,i)*w_br + a_sims_br(ind_temp2,i)*(1+r_br), eta_temp, sigma_etasq_sims(ind_temp2,i), cstar_brsim,sigma_c,psi, cov_fun);

i = 3;

ind_temp2 = 1:tt;
ind_temp2 = ind_temp2( sigma_etasq_sims(1:tt,i) < 1e10);
eta_temp = eta_sims( ind_temp2,i);
[chat_pol3,shat_pol3] =  gp_update( stateGrid' , yi(ind_temp2,i)*w_br + a_sims_br(ind_temp2,i)*(1+r_br), eta_temp, sigma_etasq_sims(ind_temp2,i), cstar_brsim,sigma_c,psi, cov_fun);


%ybar1_ind = find(chat_pol1 < stateGrid', 1,'first');
ybar1_ind = find(rw_pol< stateGrid',1,'first'); 
ybar2_ind = find(chat_pol2 > rw_pol, 1,'first'); 
ybar3_ind = find(chat_pol3 > rw_pol, 1,'first'); 

%figure
%subplot(1,2,1)
%plot(stateGrid, [chat_pol1,chat_pol2, ctilde', rw_pol])
%subplot(1,2,2)
%plot(stateGrid, [shat_pol1,shat_pol2])

chat_pol1_temp = chat_pol1; 
chat_pol1_temp( chat_pol1 > stateGrid') = stateGrid( chat_pol1 > stateGrid');


%%%%%%%%%%% Figure IV.(a)
h = figure
%if tt > 1
plot(stateGrid,chat_pol1,'-b', 'LineWidth',2)
hold on
plot(stateGrid,chat_pol2,'--r', 'LineWidth',2)

plot(stateGrid(1:18:end),chat_pol3(1:18:end),':+', 'color',[0 0.5 0], 'LineWidth',2)
plot(stateGrid,cstar_brsim(stateGrid),'-black', stateGrid, rw_pol, '-.k',stateGrid,stateGrid,':black')

scatter(stateGrid(ybar1_ind),rw_pol(ybar1_ind),180,'b','filled')
scatter(stateGrid(ybar2_ind),chat_pol2(ybar2_ind),180,'r','filled')
scatter(stateGrid(ybar3_ind),chat_pol3(ybar3_ind),180, [0 0.5 0],'filled')

xlabel('Cash on hand ($y$)','interpreter','latex')
ylabel('Consumption ($c$)','interpreter','latex')
xlim([0.6,13])
ylim([0.9,1.8])
ax = gca;
ax.TickLabelInterpreter='latex';
ax.FontSize = 14;

%legend('Agent $S$, $t=100$ ($\widehat{c}_{S,100}(y)$)','Agent $C$, $t=100$ ($\widehat{c}_{C,100}(y)$)','Agent $S$, $t=1$ ($\widehat{c}_{S,1}(y)$)','Agent $C$, $t=1$ ($\widehat{c}_{C,1}(y)$)', '$c^*(y)$','$c^{RW}(y)$','Constraint','Location','Southeast','interpreter','latex','FontSize',14, 'NumColumns',2)
legend('$\widehat{c}_{c}(y)$','$\widehat{c}_{u_1}(y)$','$\widehat{c}_{u_2}(y)$', '$c^*(y)$','$c^{RW}(y)$','Constraint','Location','Southeast','interpreter','latex','FontSize',16, 'NumColumns',2);
hold off


print('-depsc', 'chat_pol')
print('-dpng', 'chat_pol')

RGB = imread('chat_pol.png');
I = rgb2gray(RGB);
figure
imshow(I)

figure
plot( ya(1:tt,2), '-r', 'LineWidth',2)
hold on 
plot( ya(1:tt,1), '-b', 'LineWidth',2)
plot( ya(1:tt,3), '-', 'LineWidth',2)
plot( stateGrid(ybar2_ind)*ones(tt,1), ':r', 'LineWidth',1.5)
hline(w_br,':b', '$w$',14,8,1)
plot( stateGrid(ybar1_ind)*ones(tt,1), ':b', 'LineWidth',1.5)
scatter(1, ya(1,1),200,'b')
hold off

xlabel('Time')
ylabel('Cash-on-Hand')

ax = gca;
ax.TickLabelInterpreter='latex';
ax.FontSize = 13;

legend('$y_{S,t}$','$y_{C,t}$','$\bar y_{S}$', '$\bar y_C$' ,'Location','Northeast','interpreter','latex','FontSize',16, 'NumColumns',2);
print('-depsc', 'ya_evolution')

%%%%%%%%%%% Figure IV.(b)
figure
%if tt > 1
plot(stateGrid,shat_pol2,'--r', 'LineWidth',2)
hold on
plot(stateGrid,shat_pol1,'-b', 'LineWidth',2)
plot(stateGrid(1:18:end),shat_pol3(1:18:end),':x', 'color',[0 0.5 0], 'LineWidth',2)
plot(stateGrid, kappa/2*ones(size(stateGrid)), ':', 'LineWidth',2)

scatter(stateGrid(ybar1_ind),shat_pol1(ybar1_ind),180,'b','filled')
%scatter(ya(1,2),chat_sims_br(1,2),200,'r')
scatter(stateGrid(ybar2_ind),shat_pol2(ybar2_ind),180,'r','filled')
scatter(stateGrid(ybar3_ind),shat_pol3(ybar3_ind),180, [0 0.5 0],'filled')


xlabel('Cash on hand ($y$)','interpreter','latex')
ylabel('Uncertainty ($\widehat{\sigma}^2)$','interpreter','latex')
xlim([0.6,13])
%ylim([0.9,1.7])
%vline(ya(1,1),':k', '$y_{i,1}$',14,-5,1)

ax = gca;
ax.TickLabelInterpreter='latex';
ax.FontSize = 13;

%legend('Agent $S$, $t=100$ ($\widehat{c}_{S,100}(y)$)','Agent $C$, $t=100$ ($\widehat{c}_{C,100}(y)$)','Agent $S$, $t=1$ ($\widehat{c}_{S,1}(y)$)','Agent $C$, $t=1$ ($\widehat{c}_{C,1}(y)$)', '$c^*(y)$','$c^{RW}(y)$','Constraint','Location','Southeast','interpreter','latex','FontSize',14, 'NumColumns',2)
legend('$\widehat{\sigma}^2_{c}(y)$','$\widehat{\sigma}^2_{u_1}(y)$','$\widehat{\sigma}^2_{u_2}(y)$','$\kappa$','Location','Northeast','interpreter','latex','FontSize',16, 'NumColumns',2);
hold off
print('-depsc', 'shat_pol')


 